Authors Jason Codrington and Jay Foley

Authors Jason Codrington and Jay Foley#

Background#

We are going to demonstrate the use of the split-operator method to solve the time-dependent Schroedinger equation in 1D,

\[ i\hbar \frac{\partial}{\partial t} \Psi(x,t) = \hat{H} \Psi(x,t), \]

where

\[ \begin{align}\begin{aligned} \hat{H} = -\frac{\hbar^2}{2m} \frac{d^2}{dx^2} + V(x) \equiv \hat{K} + \hat{V} $$.\\The idea behind the split operator method is to approximate the time evolution operator $\hat{U}(t,t_0)$ for a very small time-increment $\Delta t$ as follows:\\$$ \hat{U}(t+\Delta t, t) = {\rm exp}\left( -\frac{i \Delta t}{2\hbar} \hat{V}\right) {\rm exp}\left( -\frac{i\Delta t}{\hbar} \hat{K} \right) {\rm exp}\left( - \frac{i \Delta t}{2\hbar} \hat{V} \right) $$, such that\\$$ \Psi(x, t+\Delta t) \approx \hat{U}(t+\Delta t, t) \Psi(x,t). \end{aligned}\end{align} \]

In this approach, we apply a half time-step update to the wavefunction through the kinetic energy operator, a full time-step update through the potential energy operator, then a second half time-step update through the kinetic energy operator.

That looks great, but how the heck do we exponentiate \(\hat{K}\) with it’s second derivatives? IDK, but actually we don’t have to. The kinetic energy operator Fourier transforms to a quadratic function in momentum space:

\[ -\frac{\hbar^2}{2m}\frac{d^2}{dx^2} \xrightarrow{\mathcal{F}} \frac{ p^2}{2m}, \]

and this can be easily exponentiated. We will denote the kinetic energy operator in momentum space as \( \hat{K}_p = \frac{p^2}{2m} \), and we can therefore write the following:

\[ {\rm exp}\left( -\frac{i \Delta t}{\hbar} \hat{K}\right) \Psi(x,t) = \mathcal{F}^{-1}\left( {\rm exp}\left( -\frac{i \Delta t}{\hbar} \hat{K}_p \right) \mathcal{F} \left( \Psi(x,t) \right) \right). \]

For our purposes, this is actually awesome because we can take the Fourier and inverse Fourier transform of things with like a single line of python code. The only mild additional discomfort we have to endure is having to carry an additional grid around - we will need a spatial grid to evaluate \(\Psi(x,t)\) and \(\hat{V}\) on, but we will need a momentum grid to evaluate \(\hat{K}_p\) on.

We can define the points on our spatial grid as follows:

\[ \begin{align}\begin{aligned} x_j = x_{min} + j \Delta x $$ for $ j = 0, 1, ..., N-1 $ where $N$ is the number of gridpoints and\\$$ \Delta x = \frac{x_{max} - x_{min}}{N-1}$$. \\Our momentum grid will be defined as follows:\\$$ p_j = j \frac{2 \pi \hbar}{N \Delta x}, \end{aligned}\end{align} \]

where \(j = -\frac{N}{2}, ..., \frac{N}{2}. \)

Outline of our code#

  1. Add attributes to our class for \(x_{min}\), \(x_{max}\), \(N\) and create grids

  2. Add attributes to our class for other parts of the simulation, including

    • time step \(\Delta t\) (in atomic units)

    • mass of particle \(m\) (in atomic units)

    • value of \(\hbar\) (in atomic units)

  3. Define initial wavefunction $\(\Psi(x,t_0) = \frac{1}{\sigma \sqrt{2\pi}} {\rm exp}\left(\frac{-1}{2}\left(\frac{x-x_0}{\sigma}\right)^2\right) {\rm exp}\left(i k_0 x \right)\)$

    • \(k_0 \approx \frac{m}{ 2}\)

    • \(x_0 \approx \frac{x_{max}}{2}\)

    • \(\sigma \approx \frac{x_{max}-x_{min}}{30} \)

  4. Define \(\hat{V}\)… could be

    • Nothing

    • Square barrier

    • Square well

    • Gaussian barrier?

    • Harmonic well?

  5. Define \(\hat{K}_p\) operator using \(\frac{{\bf p}^2}{2m}\)

  6. Form exponentials of the \(\hat{V}\) and \(\hat{K}_p\) operators

    • \({\bf K} = {\rm exp}\left( -\frac{i \Delta t}{\hbar} \hat{K}_p \right) \)

    • \({\bf V} = {\rm exp}\left( -\frac{i \Delta t}{2\hbar} \hat{V} \right) \)

  7. Implement split operator method:

  • \(\Psi^{(1)}(x,t) = {\bf V} \: \Psi(x,t) \)

  • \(\Psi^{(1)}(p,t) = \mathcal{F} \left( \Psi^{(1)}(x,t)\right)\)

  • \(\Psi^{(2)}(p,t) = {\bf K} \: \Psi^{(1)}(p,t)\)

  • \(\Psi^{(2)}(x,t) = \mathcal{F}^{-1} \left( \Psi^{(2)}(p,t)\right) \)

  • \(\Psi(x,t+\Delta t) = {\bf V} \: \Psi^{(2)}(x,t) \)

from matplotlib import animation, rc, pyplot as plt
from IPython.display import HTML
import numpy as np

class Quantum:
    def __init__(self, args):

        if 'box_length' in args:
            self.L = args['box_length']
        else:
            self.L = 10.0
        if 'time_step' in args:
            self.dt = args['time_step']
        else:
            self.dt = 0.1
        if 'grid_points' in args:
            self.grid_points =args['grid_points']
        else:
            self.grid_points = 256
        if 'mass' in args:
          self.m = args['mass']
        else:
          self.m = 1

        ''' some constants and parameters in atomic units '''
        # reduced planck's constant
        self.hbar = 1
        # reduced mass for PIR system or a very light harmonic oscillator!
        self.mu = 1
        # very small spring constant!
        self.k = 1
        # position-space spread for a gaussian wavepacket
        self.sigma = self.L / 5
        # central position value for a gaussian wavepacket
        self.x0 = self.L / 2
        # central momentum for a gaussian wavepacket
        self.k0 = 0.4
        self.x = np.linspace(0, self.L, self.grid_points)
        self.dx = self.x[1]-self.x[0]
        res = self.grid_points
        self.dk = 2 * np.pi / (res * self.dx)
        self.k = np.concatenate((np.arange(0, res / 2),
                                 np.arange(-res / 2, 0))) * self.dk

        self.V = np.zeros_like(self.x)
        #self.V = 0.5 * (self.x - self.voffset) ** 2
        self.Psi = np.sqrt(2/self.L) * np.sin((np.pi * self.x/self.L), dtype=complex)
        self.K = np.exp(-1/(2*self.m) * (self.k ** 2) * self.dt * 1j)
        self.R = np.exp(-0.5 * self.V * self.dt * 1j)






    def build_operator(self):
        self.R = np.exp(-0.5 * self.V * self.dt * 1j)

    def split_op(self):


        # Half-step in real space
        self.Psi *= self.R

        # FFT to momentum space
        self.Psi = np.fft.fft(self.Psi)

        # Full step in momentum space
        self.Psi *= self.K

        # iFFT back
        self.Psi = np.fft.ifft(self.Psi)

        # Final half-step in real space
        self.Psi *= self.R
params = {'box_length': 500, 'v_offset': 0, 'grid_points': 1000, 'wfc_offset': 0, 'system': 'pib', 'mass': 1.0,
          'time_step': 1.0}
wf = Quantum(params)
ci = 0+1j
x0 = 200
sig = 15
k0 = 0.4*1

### T1 will be the prefactor that is 1/(sigma * sqrt(2 * pi))
T1 = 1/(sig * np.sqrt(2 * np.pi))
### T2 will be the Gaussian function, exp(-0.5 * ((x-x0)/sigma)^2)
T2 = np.exp(-0.5 * ((wf.x-x0)/sig)**2)
### T3 will be the complex exponential (aka the plane wave!)
T3 = np.exp(ci * k0 * wf.x)

wf.Psi = T1 * T2 * T3

### put in a square barrier between x=600 and x=700
wf.V[600:700] = +0.025

wf.build_operator()
# First set up the figure, the axis, and the plot element we want to animate
fig, ax = plt.subplots()
plt.close()


### parameters for plot
ax.set_xlim(( 0, 500))
ax.set_ylim((-0.003, 0.003))

line, = ax.plot([], [], lw=2)
lineV, = ax.plot([], [], lw=2)

# initialization function: plot the background of each frame
def init():
    line.set_data([], [])
    lineV.set_data([], [])
    return (line, lineV,)

#def build_triangle_potential(leftpoint, midpoint, rightpoint):


def animate(i):
    wf.split_op()
    #print(i,wf.Psi[2000])
    line.set_data(wf.x, np.real(np.conj(wf.Psi)*wf.Psi))
    lineV.set_data(wf.x, wf.V)
    return(line, lineV,)


anim = animation.FuncAnimation(fig, animate, init_func=init,
                             frames=400, interval=10, blit=True)

# Note: below is the part which makes it work on Colab
rc('animation', html='jshtml')
anim
# (Optional) Utility: infer a barrier right edge from a potential array
import numpy as np

def infer_barrier_right_edge(x, V, frac=0.1):
    Vmax = np.max(V)
    thresh = frac * Vmax if Vmax > 0 else 0.0
    mask = V > thresh
    if not np.any(mask):
        return float(x[int(0.7*len(x))])
    idx = np.where(mask)[0]
    return float(x[idx.max()])
# === Feature 1: Interactive potential builder (rectangle/triangle/trapezoid) ===
import numpy as np
import ipywidgets as w
import matplotlib.pyplot as plt
from IPython.display import display, clear_output

def rectangular_potential(x, x_left, width, height):
    x_right = x_left + width
    return np.where((x >= x_left) & (x <= x_right), height, 0.0)

def triangular_potential(x, x_left, width, height, apex_pos=0.5):
    x_right = x_left + width
    V = np.zeros_like(x, dtype=float)
    apex_x = x_left + apex_pos * width
    left_mask = (x >= x_left) & (x <= apex_x)
    if np.any(left_mask):
        V[left_mask] = height * (x[left_mask] - x_left) / max(apex_x - x_left, 1e-15)
    right_mask = (x > apex_x) & (x <= x_right)
    if np.any(right_mask):
        V[right_mask] = height * (x_right - x[right_mask]) / max(x_right - apex_x, 1e-15)
    return V

def trapezoid_potential(x, x_left, width, height_left, height_right, ramp_frac=0.2):
    w_ramp = max(ramp_frac * width, 1e-15)
    x_rise_end = x_left + w_ramp
    x_fall_start = x_left + width - w_ramp
    x_right = x_left + width
    V = np.zeros_like(x, dtype=float)
    mask_rise = (x >= x_left) & (x < x_rise_end)
    if np.any(mask_rise):
        V[mask_rise] = height_left * (x[mask_rise] - x_left) / w_ramp
    mask_mid = (x >= x_rise_end) & (x <= x_fall_start)
    if np.any(mask_mid):
        mid_len = max(x_fall_start - x_rise_end, 1e-15)
        V[mask_mid] = height_left + (height_right - height_left) * (x[mask_mid] - x_rise_end) / mid_len
    mask_fall = (x > x_fall_start) & (x <= x_right)
    if np.any(mask_fall):
        V[mask_fall] = height_right * (x_right - x[mask_fall]) / w_ramp
    return V

def build_potential(x, shape, params):
    if shape == "Rectangle":
        return rectangular_potential(x, params["x_left"], params["width"], params["height"])
    elif shape == "Triangle":
        return triangular_potential(x, params["x_left"], params["width"], params["height"], params["apex_pos"])
    elif shape == "Trapezoid":
        return trapezoid_potential(x, params["x_left"], params["width"], params["height_left"], params["height_right"], params["ramp_frac"])
    else:
        return np.zeros_like(x)

shape_dd      = w.Dropdown(options=["Rectangle", "Triangle", "Trapezoid"], value="Rectangle", description="Shape:")
x_left_sl     = w.FloatSlider(value=float(x.min()) + 0.2*(x.max()-x.min()), min=float(x.min()), max=float(x.max()), step=float((x[1]-x[0])), description="x_left:", continuous_update=False)
width_sl      = w.FloatSlider(value=0.2*(x.max()-x.min()), min=float((x[1]-x[0])*2), max=float(x.max()-x.min()), step=float((x[1]-x[0])), description="Width:", continuous_update=False)
height_sl     = w.FloatSlider(value=1.0, min=0.0, max=10.0, step=0.01, description="Height:", continuous_update=False)
apex_pos_sl   = w.FloatSlider(value=0.5, min=0.0, max=1.0, step=0.01, description="Apex pos:", continuous_update=False)
height_l_sl   = w.FloatSlider(value=1.0, min=0.0, max=10.0, step=0.01, description="Left H:", continuous_update=False)
height_r_sl   = w.FloatSlider(value=0.5, min=0.0, max=10.0, step=0.01, description="Right H:", continuous_update=False)
ramp_frac_sl  = w.FloatSlider(value=0.2, min=0.0, max=0.45, step=0.01, description="Ramp frac:", continuous_update=False)

pot_out = w.Output()
V_current = {"V": np.zeros_like(x)}

def update_plot(*_):
    with pot_out:
        clear_output(wait=True)
        shape = shape_dd.value
        params = dict(
            x_left=x_left_sl.value,
            width=width_sl.value,
            height=height_sl.value,
            apex_pos=apex_pos_sl.value,
            height_left=height_l_sl.value,
            height_right=height_r_sl.value,
            ramp_frac=ramp_frac_sl.value
        )
        V = build_potential(x, shape, params)
        V_current["V"] = V
        fig, ax = plt.subplots(figsize=(6,3))
        ax.plot(x, V)
        ax.set_title(f"{shape} potential")
        ax.set_xlabel("x")
        ax.set_ylabel("V(x)")
        ax.grid(True, alpha=0.3)
        plt.show()

def _toggle_controls(*_):
    shape = shape_dd.value
    height_sl.layout.display   = "flex" if shape == "Rectangle" else "none"
    apex_pos_sl.layout.display = "flex" if shape == "Triangle"  else "none"
    height_l_sl.layout.display = "flex" if shape == "Trapezoid" else "none"
    height_r_sl.layout.display = "flex" if shape == "Trapezoid" else "none"
    ramp_frac_sl.layout.display= "flex" if shape == "Trapezoid" else "none"
    update_plot()

for wid in [shape_dd, x_left_sl, width_sl, height_sl, apex_pos_sl, height_l_sl, height_r_sl, ramp_frac_sl]:
    wid.observe(update_plot, names="value")
shape_dd.observe(_toggle_controls, names="value")
_toggle_controls()

controls = w.VBox([shape_dd, x_left_sl, width_sl, height_sl, apex_pos_sl, height_l_sl, height_r_sl, ramp_frac_sl])
ui = w.HBox([controls, pot_out])
display(ui)

# To use in your solver: replace your V with V_current["V"] before the time loop.
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[5], line 52
     49         return np.zeros_like(x)
     51 shape_dd      = w.Dropdown(options=["Rectangle", "Triangle", "Trapezoid"], value="Rectangle", description="Shape:")
---> 52 x_left_sl     = w.FloatSlider(value=float(x.min()) + 0.2*(x.max()-x.min()), min=float(x.min()), max=float(x.max()), step=float((x[1]-x[0])), description="x_left:", continuous_update=False)
     53 width_sl      = w.FloatSlider(value=0.2*(x.max()-x.min()), min=float((x[1]-x[0])*2), max=float(x.max()-x.min()), step=float((x[1]-x[0])), description="Width:", continuous_update=False)
     54 height_sl     = w.FloatSlider(value=1.0, min=0.0, max=10.0, step=0.01, description="Height:", continuous_update=False)

NameError: name 'x' is not defined
# === Feature 2: Transmitted probability via integrated probability current ===
import numpy as np

def probability_current_1d(psi, dx, hbar=1.0, m=1.0):
    dpsi_dx = np.zeros_like(psi, dtype=complex)
    dpsi_dx[1:-1] = (psi[2:] - psi[:-2]) / (2*dx)
    dpsi_dx[0]    = (psi[1] - psi[0]) / (dx)
    dpsi_dx[-1]   = (psi[-1] - psi[-2]) / (dx)
    j = (hbar/m) * np.imag(np.conj(psi) * dpsi_dx)
    return j

class FluxDetector:
    def __init__(self, x, x_det_index, dt, hbar=1.0, m=1.0, smooth=0):
        self.x = x
        self.idx = int(x_det_index)
        self.dt = float(dt)
        self.hbar = float(hbar)
        self.m = float(m)
        self.accum = 0.0
        self._buffer = []
        self._smooth = int(smooth)

    def update(self, psi):
        dx = self.x[1] - self.x[0]
        j = probability_current_1d(psi, dx, self.hbar, self.m)
        val = float(j[self.idx])
        if self._smooth > 0:
            self._buffer.append(val)
            if len(self._buffer) > self._smooth:
                self._buffer.pop(0)
            val = sum(self._buffer)/len(self._buffer)
        self.accum += val * self.dt

    @property
    def P_transmitted(self):
        return float(self.accum)

# --- Quick start ---
# x_d = infer_barrier_right_edge(x, V) + 3*(x[1]-x[0])  # small buffer
# x_det_index = int(np.argmin(np.abs(x - x_d)))
# det = FluxDetector(x, x_det_index, dt, hbar=hbar, m=m, smooth=2)
# In your time loop (after updating K each step):  det.update(K)
# After the loop:  print("Transmitted probability (flux):", det.P_transmitted)